/*++
Copyright (c) Advanced Windows Debugging (ISBN 0321374460) from Addison-Wesley Professional.  All rights reserved.

    THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF ANY
    KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
    IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A PARTICULAR
    PURPOSE.

--*/


#include <windows.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <iostream>

#define TRACE(a) wprintf(a); wprintf  (L"Last error = %d\n", GetLastError());
#define TRACE_ERR(a,err) wprintf(a); wprintf  (L"Last error = %d\n", err);

HANDLE targetProcessHandle = NULL; 

WCHAR * ReadRemoteString(HANDLE process, LPVOID address, WORD length, WORD unicode)
{
    WCHAR * msg = new WCHAR[length];
    if (!msg) return NULL;
    memset(msg, 0, sizeof(WCHAR)*(length));

    if ( unicode )
    {
        ReadProcessMemory(process, address ,msg, length*sizeof(WCHAR),NULL);  
        return msg;
    } 
    else
    {
        BYTE * originalMsg = new BYTE[length];
        if (!originalMsg)
        {
            delete[] msg;
            return NULL;
        }
        memset(originalMsg, 0, sizeof(BYTE)*(length));
        
        ReadProcessMemory(process, address ,originalMsg, length,NULL);    
        for (WORD i = 1; i < length; i++)
        {
            msg[i] = originalMsg[i];        
        }
        return msg;
    }
}

WCHAR ReadCharW(HANDLE process, LPVOID address)
{
    WCHAR result = '\0';
    ReadProcessMemory(process, address ,&result, sizeof(WCHAR),NULL);
    return result;
};

WCHAR ReadCharA(HANDLE process, LPVOID address)
{
    WCHAR result = '\0';
    ReadProcessMemory(process, address ,&result, sizeof(CHAR),NULL);
    return result;
};

WCHAR * ReadRemoteSZ(HANDLE process, LPVOID address, WORD unicode)
{
    if (NULL == address) return NULL;
    WCHAR * msg = 0 ;
    if ( unicode )
    {
        WCHAR * ptr = (WCHAR*)address;
        size_t stringLen = 0;
        while(ReadCharW(process, ptr++)) stringLen++; 
        msg = new WCHAR[stringLen+1];
        if (!msg) return NULL;
        memset(msg, 0, sizeof(WCHAR)*(stringLen+1));
        ReadProcessMemory(process, address ,msg, (stringLen+1)*sizeof(WCHAR),NULL);  
        return msg;
    } 
    else
    {
        CHAR * ptr = (CHAR*)address;
        size_t stringLen = 0;
        while(ReadCharA(process, ptr++)) stringLen++; 

        BYTE * originalMsg = new BYTE[stringLen+1];
        if (!originalMsg) return NULL;
        memset(originalMsg, 0, sizeof(BYTE)*(stringLen+1));
        
        msg = new WCHAR[stringLen]; 
        if (!msg) 
        {
            delete[] originalMsg;
            return NULL;
        }
        memset(msg, 0, sizeof(WCHAR)*(stringLen+1));

        ReadProcessMemory(process, address ,originalMsg, stringLen+1,NULL);    
        for (size_t i = 1; i < stringLen+1; i++)
        {
            msg[i] = originalMsg[i];        
        }
        delete[] originalMsg;
        return msg;
    }
}

VOID * ReadRemotePtr(HANDLE process, LPVOID address)
{
    if (NULL == address) return NULL;
    VOID* newPointer = 0 ;
    if (!ReadProcessMemory(process, address ,&newPointer, sizeof(VOID *),NULL))
    {
        newPointer=NULL;
    };  
    return newPointer;
}

DWORD ProcessEvent(DEBUG_EVENT& dbgEvent)
{
    std::cout << "New event generated by PID.TID=" << dbgEvent.dwProcessId <<"." << dbgEvent.dwThreadId << std::endl;

    switch (dbgEvent.dwDebugEventCode)
    {
        case EXCEPTION_DEBUG_EVENT:

            //typedef struct _EXCEPTION_DEBUG_INFO {
            //    EXCEPTION_RECORD ExceptionRecord;
            //    DWORD dwFirstChance;
            //} EXCEPTION_DEBUG_INFO;

            std::cout << "\tExceptionEvent:\t\n\t Code = " << std::hex << 
            dbgEvent.u.Exception.ExceptionRecord.ExceptionCode <<  "\n\t FirstChance = " << 
            dbgEvent.u.Exception.dwFirstChance <<std::endl;

            switch (dbgEvent.u.Exception.ExceptionRecord.ExceptionCode)
            {
                case EXCEPTION_BREAKPOINT:
                case EXCEPTION_SINGLE_STEP:
                case 0x4000001f:
                case 0x4000001e:
                    return DBG_CONTINUE;
            }
            return DBG_EXCEPTION_NOT_HANDLED;

        case CREATE_THREAD_DEBUG_EVENT:

            //typedef struct _CREATE_THREAD_DEBUG_INFO {
            //    HANDLE hThread;
            //    LPVOID lpThreadLocalBase;
            //    LPTHREAD_START_ROUTINE lpStartAddress;
            //} CREATE_THREAD_DEBUG_INFO;
            std::cout << "\tCreateThreadEvent TID=" << dbgEvent.dwThreadId <<std::endl;
            break;

        case CREATE_PROCESS_DEBUG_EVENT:
            //typedef struct _CREATE_PROCESS_DEBUG_INFO {
            //    HANDLE hFile;
            //    HANDLE hProcess;
            //    HANDLE hThread;
            //    LPVOID lpBaseOfImage;
            //    DWORD dwDebugInfoFileOffset;
            //    DWORD nDebugInfoSize;
            //    LPVOID lpThreadLocalBase;
            //    LPTHREAD_START_ROUTINE lpStartAddress;
            //    LPVOID lpImageName;
            //    WORD fUnicode;
            //} CREATE_PROCESS_DEBUG_INFO, *LPCREATE_PROCESS_DEBUG_INFO;
            {
            targetProcessHandle = dbgEvent.u.CreateProcessInfo.hProcess;
            std::wcout << L"\tCreateProcessEvent PID=" << dbgEvent.dwProcessId << std::endl;
            }
            break;

        case EXIT_THREAD_DEBUG_EVENT:
            //typedef struct _EXIT_THREAD_DEBUG_INFO {
            //    DWORD dwExitCode;
            //} EXIT_THREAD_DEBUG_INFO, *LPEXIT_THREAD_DEBUG_INFO;
            std::cout << "\tExitThreadEvent TID=" << dbgEvent.dwThreadId << " ExitCode=" << 
            dbgEvent.u.ExitThread.dwExitCode<<std::endl;
            break;

        case EXIT_PROCESS_DEBUG_EVENT:
            //typedef struct _EXIT_PROCESS_DEBUG_INFO {
            //    DWORD dwExitCode;
            //} EXIT_PROCESS_DEBUG_INFO, *LPEXIT_PROCESS_DEBUG_INFO;
            std::cout << "\tExitProcessEvent ExitCode=" << std::hex << dbgEvent.u.ExitProcess.dwExitCode <<std::endl;
            return 0 ;
            break;

        case LOAD_DLL_DEBUG_EVENT:
            //typedef struct _LOAD_DLL_DEBUG_INFO {
            //    HANDLE hFile;
            //    LPVOID lpBaseOfDll;
            //    DWORD dwDebugInfoFileOffset;
            //    DWORD nDebugInfoSize;
            //    LPVOID lpImageName;
            //    WORD fUnicode;
            //} LOAD_DLL_DEBUG_INFO, *LPLOAD_DLL_DEBUG_INFO;
            {    
                std::cout << "\tLoadDllEvent ";
                LPVOID imageAddress = ReadRemotePtr(targetProcessHandle, dbgEvent.u.LoadDll.lpImageName);
                WCHAR * msg = ReadRemoteSZ(targetProcessHandle, imageAddress, dbgEvent.u.LoadDll.fUnicode);

                if (NULL == msg)
                {
                    std::wcout << L"ImageName: cannot be read"<< std::endl;
                }
                else
                {
                    std::wcout << L"ImageName: " << msg << std::endl;
                    delete[] msg;
                }
            }
            break;

        case UNLOAD_DLL_DEBUG_EVENT:
            
            //typedef struct _UNLOAD_DLL_DEBUG_INFO {
            //    LPVOID lpBaseOfDll;
            //} UNLOAD_DLL_DEBUG_INFO, *LPUNLOAD_DLL_DEBUG_INFO;

            std::cout << "\tUnloadDllEvent Address=" << std::hex << dbgEvent.u.UnloadDll.lpBaseOfDll <<std::endl;
            break;

        case OUTPUT_DEBUG_STRING_EVENT:

            //typedef struct _OUTPUT_DEBUG_STRING_INFO {
            //    LPSTR lpDebugStringData;
            //    WORD fUnicode;
            //    WORD nDebugStringLength;
            //} OUTPUT_DEBUG_STRING_INFO, *LPOUTPUT_DEBUG_STRING_INFO;

            {
                std::cout << "\tOutputDebugEvent ";
                OUTPUT_DEBUG_STRING_INFO& OutputDebug = dbgEvent.u.DebugString;
                WCHAR * msg = ReadRemoteString(targetProcessHandle, OutputDebug.lpDebugStringData, OutputDebug.nDebugStringLength, OutputDebug.fUnicode);
                if (NULL == msg)
                {
                    std::wcout << L"Cannot read OutputDebug"<< std::endl;
                }
                else
                {
                    std::wcout << L"OutputDebug:\t"<< msg << std::endl;
                    delete[] msg;
                }
                break;
            }

        case RIP_EVENT:

            //typedef struct _RIP_INFO {
            //    DWORD dwError;
            //    DWORD dwType;
            //} RIP_INFO, *LPRIP_INFO;

            std::cout << "RipEvent Error="<< std::hex << dbgEvent.u.RipInfo.dwError <<std::endl;
            break;
    }
    return DBG_CONTINUE ;
}

VOID _cdecl wmain( ULONG argc, WCHAR* argv[] )
{
    if (argc < 2)
    {
        wprintf(L"03sample <target>\n");
        return;
    }

    STARTUPINFO startupInfo={0}; 
    startupInfo.cb = sizeof(startupInfo);
    PROCESS_INFORMATION processInfo = {0};


    int commandLength = wcslen(argv[1])+(argv[2]?wcslen(argv[2]):0) + 2;
    WCHAR * commandLine = new WCHAR[commandLength];
    if (!commandLine)
    {
        TRACE(L"Allocation failed\n");
    }
    wcscpy_s(commandLine, commandLength, argv[1]);
    if (argv[2])
    {
        wcscat_s(commandLine, commandLength, L" ");
        wcscat_s(commandLine, commandLength, argv[2]);
    }

    BOOL res = CreateProcess(NULL, commandLine, NULL, NULL, FALSE, DEBUG_PROCESS, NULL, NULL, &startupInfo, &processInfo);
    if (FALSE == res)
    {
        TRACE(L"CreateProcess failed\n");
        return;
    }
    DEBUG_EVENT debugEvent = { 0 } ;

    DWORD endDisposition = DBG_CONTINUE;
    for(;endDisposition != 0;)
    {
        if (!WaitForDebugEvent(&debugEvent, INFINITE))
        {
            TRACE(L"WaitForDebugEvent failed\n");
            break;
        }
        endDisposition = ProcessEvent(debugEvent);
        if (0 == endDisposition) break;
        if (!ContinueDebugEvent(debugEvent.dwProcessId, debugEvent.dwThreadId, endDisposition))
        {
            TRACE(L"ContinueDebugEvent failed\n");
            break;
        };
    }
}

